#include "faceauthenticate_cv.h"

#ifdef USE_OPENCV


#include <src/utils/log.h>


static const int g_face_label_id = 10010001;

FaceAuthenticateCV::FaceAuthenticateCV(QObject *parent)
    : FaceAuthInterface{parent}
{

}

FaceAuthenticateCV::~FaceAuthenticateCV()
{

}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::init()
{


    return FaceAuthInterface::FaceAuth_OK;
}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::finit()
{

    return FaceAuthInterface::FaceAuth_OK;
}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::face_detect(void* datainput, void* dataoutput)
{
    if (m_cascade.empty())
        return FaceAuth_Init_Failed;

    Mat* pSrcImage = static_cast<Mat*>(datainput);
    Mat grayImage, dstImage;

    dstImage = pSrcImage->clone();

    grayImage.create(pSrcImage->size(), pSrcImage->type());
    cvtColor(*pSrcImage, grayImage, COLOR_BGR2GRAY);

    vector<Rect> rect;
    m_cascade.detectMultiScale(grayImage, rect, 1.1, 3, 0);

    if (rect.size() > 1)
        return FaceAuth_Failed_Multiface;
    if (rect.empty())
        return FaceAuth_Failed_Noface;

    Rect* distRect = static_cast<Rect*>(dataoutput);
    distRect->x = rect[0].x;
    distRect->y = rect[0].y;
    distRect->width = rect[0].width;
    distRect->height = rect[0].height;

    rectangle(*pSrcImage, rect[0], CV_RGB(255,0,0), 3);
    return FaceAuthInterface::FaceAuth_OK;
}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::face_regist(const QString key, void* datainput)
{

    return FaceAuthInterface::FaceAuth_OK;
}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::face_predict(void* datainput, FaceRecogResult& result)
{

    return FaceAuthInterface::FaceAuth_OK;
}


void FaceAuthenticateCV::initFromModelPath(const QList<QString>& model_paths, const QString& detect_path)
{
    int len = model_paths.size();
    m_models = vector<CVModel>(len);
    m_face_model_paths = vector<String>(len);

    FILE *file = nullptr;
    for (int i=0;i<len;i++)
    {
        String face_model_path = model_paths[i].toStdString();
        m_face_model_paths[i] = face_model_path;
        file = fopen(face_model_path.c_str(), "r");
        if (file)
        {
            fclose(file);
            m_models[i] = Algorithm::load<LBPHFaceRecognizer>(face_model_path);
//            INFO("init model from " << model_paths[i] << m_models[i].empty());
        }
        else
        {
            m_models[i].reset();
//            WARNING("init model path not vaild: " << model_paths[i]);
        }
    }

    m_face_detect_path = detect_path.toStdString();
    file = fopen(m_face_detect_path.c_str(), "r");
    if (file)
    {
        fclose(file);
        m_cascade.load(m_face_detect_path);
//        INFO("init detect from " << detect_path << m_cascade.empty());
    }
    else
    {
//        WARNING("init detect path not vaild: " << detect_path);
    }
}

QList<QString> FaceAuthenticateCV::getLabels()
{
    QSet<QString> result;
    for (auto model : m_models)
    {
        if (model.empty())
            continue;
        vector<int> alllabels = model->getLabels();

        for (auto label : alllabels)
        {
            QString labelOne = QString::fromStdString(model->getLabelInfo(label));
            if (labelOne.isEmpty() || labelOne.simplified().isEmpty())
                continue;
            result.insert(labelOne);
        }
    }

    QList<QString> res = QList<QString>::fromSet(result);
    return res;
}

bool FaceAuthenticateCV::isLableExists(const QString& label)
{
    for (auto model : m_models)
    {
        if (model.empty())
            continue;

        vector<int> labels_exists = model->getLabelsByString(label.toStdString());
        if (!labels_exists.empty())
            return true;
    }
    return false;
}

void FaceAuthenticateCV::removeFaceLabel(const QString& label)
{
    int len = m_models.size();
    for (int i=0; i<len; i++)
    {
        if (m_models[i].empty())
            continue;

        vector<int> labels_exists = m_models[i]->getLabelsByString(label.toStdString());
        if (!labels_exists.empty())
        {
            m_models[i]->clear();
            m_models[i].release();
            // remove model file
            remove(m_face_model_paths[i].c_str());
        }
    }
}

bool FaceAuthenticateCV::editFaceLabel(const QString& label, const QString& newLabel)
{

    if (!isLableExists(label))
        return false;

    if (isLableExists(newLabel))
        return false;

    int len = m_models.size();
    for (int i=0; i<len; i++)
    {
        if (m_models[i].empty())
            continue;

        vector<int> labels_exists = m_models[i]->getLabelsByString(label.toStdString());
        if (labels_exists.empty())
            continue;

        m_models[i]->setLabelInfo(labels_exists[0], newLabel.toStdString());
        m_models[i]->save(m_face_model_paths[i]);
        return true;
    }
    return false;
}


FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::face_detect_ex(Mat& mat_src, Mat& mat_gray, Mat& mat_withrect, Rect& distRect)
{
    if (m_cascade.empty())
        return FaceAuth_Init_Failed;

    mat_withrect = mat_src.clone();
    mat_gray.create(mat_src.size(), mat_src.type());
    cvtColor(mat_src, mat_gray, COLOR_BGR2GRAY);

    vector<Rect> rect;
    m_cascade.detectMultiScale(mat_gray, rect, 1.1, 3, 0);

    if (rect.size() > 1)
        return FaceAuth_Failed_Multiface;
    if (rect.empty())
        return FaceAuth_Failed_Noface;

    distRect.x = rect[0].x;
    distRect.y = rect[0].y;
    distRect.width = rect[0].width;
    distRect.height = rect[0].height;

    rectangle(mat_withrect, rect[0], CV_RGB(255,0,0), 2);
    return FaceAuthInterface::FaceAuth_OK;
}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::face_predict_ex(Mat& mat_gray, FaceRecogResult& result)
{
    if (mat_gray.empty())
        return FaceAuth_Failed_InvalidImage;

    int predict_index = 0;
    int len = m_models.size();
    vector<pair<int,double>> all_confidence(len);

    for (int i=0;i<len;i++)
    {
        int predicted_label = -1;
        double predicted_confidence = 0;
        if (m_models[i].empty())
        {
            all_confidence[i] = make_pair(i, 999);
            continue;
        }
        m_models[i]->predict(mat_gray, predicted_label, predicted_confidence);
        all_confidence[i] = make_pair(i, predicted_confidence);
    }

    sort(all_confidence.begin(), all_confidence.end(), [](pair<int, double> p1, pair<int, double> p2){
        return p1.second < p2.second;
    });
    predict_index = all_confidence[0].first;
    if (m_models[predict_index].empty())
    {
        WARNING("no model for predict " <<QString::fromStdString(m_face_model_paths[predict_index]));
        result.confidence = 999;
        result.key = "";
        return FaceAuth_OK;
    }
    vector<int> labels_exists = m_models[predict_index]->getLabels();
    if (labels_exists.empty())
    {
        CRITICAL("invalid model no labels " <<QString::fromStdString(m_face_model_paths[predict_index]));
        result.confidence = 999;
        result.key = "";
        return FaceAuth_OK;
    }
    result.confidence = all_confidence[0].second;
    result.key = QString::fromStdString(m_models[predict_index]->getLabelInfo(labels_exists[0]));
    return FaceAuth_OK;
}

FaceAuthInterface::FaceAuthResult FaceAuthenticateCV::face_regist_images(const QString& label, std::vector<cv::Mat>& mats)
{
    if (label.isEmpty() || mats.empty())
        return FaceAuth_Failed_InvalidImage;

    int len = m_models.size();
    if (isLableExists(label))
    {
        for (int i=0;i<len;i++)
        {
            if (m_models[i].empty())
                continue;
            vector<int> labels_exists = m_models[i]->getLabelsByString(label.toStdString());
            if (labels_exists.empty())
                continue;

            INFO("update model for " << label << QString::fromStdString(m_face_model_paths[i]));
            vector<int> newlabels(mats.size(), labels_exists[0]);
            m_models[i]->update(mats, newlabels);
            m_models[i]->setLabelInfo( labels_exists[0], label.toStdString());
            m_models[i]->save(m_face_model_paths[i]);
            break;
        }
    }
    else
    {
        for (int i=0;i<len;i++)
        {
            if (!m_models[i].empty())
                continue;
            INFO("new model for " << label << QString::fromStdString(m_face_model_paths[i]));
            m_models[i] = LBPHFaceRecognizer::create();
            vector<int> labels(mats.size(), g_face_label_id);
            m_models[i]->train(mats,labels);
            m_models[i]->setLabelInfo(g_face_label_id, label.toStdString());
            m_models[i]->save(m_face_model_paths[i]);
            break;
        }
    }
    return FaceAuth_OK;
}

#endif // USE_OPENCV

